#include "Drrt.h"

#include "Graph.h"
#include "Sampler.h"
#include "prm_2d.h"

const Point_2 EMPTY_POINT(-7149, -8200);

Drrt::Drrt(const Prm2D& robot1Graph,
           const Prm2D& robot2Graph,
           const CollisionDetector2D& robotCollisionDetection,
           const CollisionDetector& compositeCollisionDetector,
           const Sampler& compositeSampler,
           const PointPair& start, 
           const PointPair& target)
    : m_robot1Prm(&robot1Graph), m_robot2Prm(&robot2Graph),
      m_robotCollisionDetection(&robotCollisionDetection),
      m_compositeCollisionDetection(&compositeCollisionDetector),
      m_compositeSampler(&compositeSampler),
      m_start(&start), m_target(&target) { }

void Drrt::runLoop(int numIterations) {
	m_startGraph.add_vertex(*m_start);
  m_startKd.insert(pointDFromPointPair(*m_start));
	m_targetGraph.add_vertex(*m_target);
  m_targetKd.insert(pointDFromPointPair(*m_target));

	for (int i=0; i < numIterations; i++) {
		//Expand trees
		expand(m_startGraph, m_startKd);
		expand(m_targetGraph, m_targetKd);
		vector<PointPair> treeConnection = connectTrees();
		if (treeConnection.size() > 0) {
			m_bestPath = generatePath(treeConnection);
			cout << "SOLUTION!" << endl;
			return;
		}
	}
}

void Drrt::expand(CompositeGraph& graph, Kd_tree_d<Kernel_d>& kdTree) {
	static int N = 20;
	for (int i=0; i<N; i++) {
		Point_d qRand = m_compositeSampler->generate_sample();
		PointPair qNear = nearestNeighbor(kdTree, qRand);
    PointPair qNew;
    for (int j = 0; j < 10; ++j) {
      qNew = compositeDirectionOracle(graph, qRand, qNear);
      if (qNew.first != EMPTY_POINT) {
        break;
      }
    }
    if (qNew.first == EMPTY_POINT) {
      continue;
    }
		if (!graph.is_in_graph(qNew)) {
			graph.add_vertex(qNew);
      kdTree.insert(pointDFromPointPair(qNew));
			graph.add_edge(qNear, qNew, configurationDistance(qNear, qNew));
		}
	}
}

vector<PointPair> Drrt::connectTrees() {
	vector<PointPair> retVal;

	static int P = 5;
	for (int sample = 0; sample < P; sample++) {
		Point_d qRand = m_compositeSampler->generate_sample();
		PointPair q0 = nearestNeighbor(m_startKd, qRand);
		PointPair q1 = nearestNeighbor(m_targetKd, qRand);

    // We don't want to try to connect trees unless they're sufficiently close.
    if (configurationDistance(q0, q1) > 75) {
      continue;
    }

    vector<PointPair> connection = localConnector(q0, q1);
    if (connection.size() != 0) {
      retVal.insert(retVal.end(), connection.begin(), connection.end());
      break;
    }
	}
	return retVal;
}

class ComparePointPair {
  public:
    ComparePointPair(map<PointPair, double>& fScore) : m_fScore(fScore) { }
    bool operator()(const PointPair& a, const PointPair& b) const {
      return m_fScore[a] >= m_fScore[b];
    }
  private:
    map<PointPair, double>& m_fScore;
};

vector<PointPair> Drrt::localConnector(PointPair start, PointPair goal) const {
  set<PointPair> closedSet;
  map<PointPair, PointPair> cameFrom;
  map<PointPair, double> gScore;
  map<PointPair, double> fScore;
  gScore[start] = 0;
  fScore[start] = configurationDistance(start, goal);
  ComparePointPair fScoreComparator(fScore);
  priority_queue<PointPair, vector<PointPair>, ComparePointPair> openQueue(fScoreComparator);
  set<PointPair> openSet;
  openQueue.push(start);
  openSet.insert(start);

  int iters = 0;
  while ((openQueue.size() > 0) && (iters < 5)) {
    iters++;
    PointPair current = openQueue.top();
    if (current == goal) {
      return reconstructPath(cameFrom, start, goal);
    }

    openQueue.pop();
    openSet.erase(current);
    closedSet.insert(current);
    vector<PointPair> neighbors = getNeighbors(current);
    for (vector<PointPair>::const_iterator neighbor = neighbors.begin(); neighbor != neighbors.end(); ++neighbor) {
      double tentativeGScore = gScore[current] + configurationDistance(current, *neighbor);
      if ((closedSet.find(*neighbor) != closedSet.end()) && (tentativeGScore >= gScore[*neighbor])) {
        continue;
      }

      if ((openSet.find(*neighbor) == openSet.end()) || (tentativeGScore < gScore[*neighbor])) {
        cameFrom[*neighbor] = current;
        gScore[*neighbor] = tentativeGScore;
        fScore[*neighbor] = tentativeGScore + configurationDistance(*neighbor, goal);
        openSet.insert(*neighbor);
        openQueue.push(*neighbor);
      }
    }
  }
  return vector<PointPair>();
}

Point_2 deltaBetweenRobots(PointPair pointPair) {
	Vector_2 distanceVector = (pointPair.second - pointPair.first);
	Point_2 distancePoint = Point_2(distanceVector.x(), distanceVector.y());
	return distancePoint;
}

vector<PointPair> Drrt::reconstructPath(map<PointPair, PointPair> cameFrom, PointPair start, PointPair goal) const {
  vector<PointPair> retVal;
  retVal.push_back(goal);
  PointPair current = goal;
  while (current != start) {
    current = cameFrom[current];
    retVal.push_back(current);
  }
  reverse(retVal.begin(), retVal.end());
  return retVal;
}

vector<PointPair> Drrt::getNeighbors(PointPair current) const {
  vector<PointPair> retVal;
  vector<Point_2> robot1Neighbors = m_robot1Prm->getGraph().get_neighbors(current.first);
  vector<Point_2> robot2Neighbors = m_robot2Prm->getGraph().get_neighbors(current.second);
  for (size_t i = 0; i < robot1Neighbors.size(); ++i) {
    for (size_t j = 0; j < robot2Neighbors.size(); ++j) {
      PointPair possibleNeighbor = PointPair(robot1Neighbors[i], robot2Neighbors[j]);
      if (m_robotCollisionDetection->local_planner(
            deltaBetweenRobots(current), deltaBetweenRobots(possibleNeighbor), 0.1)) {
        retVal.push_back(possibleNeighbor);
      }
    }
  }
  return retVal;
}

Point_2 Drrt::directionOracle(const Graph<Point_2>& graph, Point_2 qNear, Point_2 qRand) {
  double bestScore = 0;
  Point_2 bestVertex = EMPTY_POINT;
  vector<Point_2> neigbors = graph.get_neighbors(qNear);
  for (size_t i = 0; i < neigbors.size(); ++i) {
    Vector_2 v1 = neigbors[i] - qNear;
    v1 = v1 / std::sqrt(CGAL::to_double(v1 * v1));
    Vector_2 v2 = qRand - neigbors[i];
    v2 = v2 / std::sqrt(CGAL::to_double(v2 * v2));
    double score = CGAL::to_double(v1 * v2);
    if (score > bestScore) {
      bestScore = score;
      bestVertex = neigbors[i];
    }
  }
  return bestVertex;
}

PointPair Drrt::compositeDirectionOracle(
    const CompositeGraph& graph, const Point_d& qRand, const PointPair& qNear) {
  Point_2 robot1Rand = Point_2(qRand[0], qRand[1]);
  Point_2 robot2Rand = Point_2(qRand[2], qRand[3]);
  Point_2 robot1Candidate = directionOracle(m_robot1Prm->getGraph(), qNear.first, robot1Rand);
  Point_2 robot2Candidate = directionOracle(m_robot2Prm->getGraph(), qNear.second, robot2Rand);
  PointPair candidate = PointPair(robot1Candidate, robot2Candidate);
  if (m_robotCollisionDetection->local_planner(
	  deltaBetweenRobots(qNear), deltaBetweenRobots(candidate), 0.1)) {
    return candidate;
  }
  return PointPair(EMPTY_POINT, EMPTY_POINT);
}


vector<PointPair> Drrt::retrieve_path() {
	return m_bestPath;
}

PointPair Drrt::nearestNeighbor(Kd_tree_d<Kernel_d>& kdTree, Point_d point) {
  return pointPairFromPointD(kdTree.nearest_neighbor(point));
}

PointPair Drrt::pointPairFromPointD(const Point_d& point) const {
  Point_2 a = Point_2(point[0], point[1]);
  Point_2 b = Point_2(point[2], point[3]);
  return PointPair(a, b);
}

Point_d Drrt::pointDFromPointPair(const PointPair& pp) const {
  vector<double> v(4);
  v[0] = CGAL::to_double(pp.first[0]);
  v[1] = CGAL::to_double(pp.first[1]);
  v[2] = CGAL::to_double(pp.second[0]);
  v[3] = CGAL::to_double(pp.second[1]);
  return Point_d(v.size(), v.begin(), v.end());
}

double Drrt::configurationDistance(const PointPair& a, const PointPair& b) const {
	return m_compositeCollisionDetection->cost_of_path(pointDFromPointPair(a), pointDFromPointPair(b));
}

vector<PointPair> Drrt::generatePath(vector<PointPair>& treeConnection) {
  std::list<PointPair> fromStart;
  m_startGraph.find_path(*m_start, treeConnection[0], fromStart);

  std::list<PointPair> toTarget;
  m_targetGraph.find_path(treeConnection[treeConnection.size()-1], *m_target, toTarget);

  vector<PointPair> retVal;
  std::copy(fromStart.begin(), fromStart.end(), std::back_inserter(retVal));
  std::copy(treeConnection.begin(), treeConnection.end(), std::back_inserter(retVal));
  std::copy(toTarget.begin(), toTarget.end(), std::back_inserter(retVal));

  //Remove adjacent vertices
  vector<PointPair> retValNoCopies;
  retValNoCopies.push_back(retVal[0]);
  //cout << 0 << ":" << retVal[0].first << "," << retVal[0].second << endl;
  for (size_t i=1; i<retVal.size(); i++) {
	  double dist = configurationDistance(retVal[i-1], retVal[i]);
	  if (dist > 0.01) {
		  retValNoCopies.push_back(retVal[i]);
		  //cout << i << ":" << retVal[i].first << "," << retVal[i].second << endl;
	  }	  
  }
  
  return retValNoCopies;
}
